//
// Created by zf on 2025/4/21.
//
#include <iostream>
#include <memory>
#include <atomic>
#include <thread>
#include <vector>
#include <cassert>
#include <mutex>
#include<algorithm>

using namespace std;


template<typename T>
class LockFreeQueue {
private:
    struct Node{
        std::shared_ptr<T> data;
        Node * next;
        Node():next(nullptr){}
    };
private:
    std::atomic<Node *> head;
    std::atomic<Node *> tail;
public:
    LockFreeQueue(){
        Node * node = new Node();
        head.store(node);
        tail.store(node);
    }
    ~LockFreeQueue(){
        while(Node * const curr = head.load()){
            head.store(curr->next);
            delete curr;
        }
    }

    LockFreeQueue(const LockFreeQueue &) = delete;
    LockFreeQueue & operator = (const LockFreeQueue &) = delete;

    void push(T val){
        Node * node = new Node();
        std::shared_ptr<T> data = std::make_shared<T>(std::move(val));
        Node * const old_tail = tail.load();
        old_tail->data.swap(data);
        old_tail->next = node;
        tail.store(node);
    }
    std::shared_ptr<T> pop(){
        Node * node = pop_head();
        if(!node){
            return std::shared_ptr<T>();
        }
        head.store(node->next);
        std::shared_ptr<T> res (node->data);
        delete node;
        return res;
    }

    bool empty(){
        return head.load() == tail.load();
    }
private:
    Node * pop_head(){
        Node * const old_head = head.load();
        if(old_head == tail.load()){
            return nullptr;
        }
        head.store(old_head->next);
        return old_head;
    }

};






template<typename T>
class MPMCLockFreeQueue {
private:
    struct Node;

    struct alignas(16) CountedNodePtr {
        int external_count;
        Node* ptr;
    };

    struct Node {
        std::atomic<T*> data;
        std::atomic<int> internal_count;
        std::atomic<CountedNodePtr> next;

        Node() : data(nullptr), internal_count(0), next({0, nullptr}) {}

        void release_ref() {
            // 减少内部计数，如果归零则删除节点
            int old_count = internal_count.fetch_sub(1, std::memory_order_acq_rel);
            if (old_count == 1) {
                delete this;
            }
        }
    };

    std::atomic<CountedNodePtr> head;
    std::atomic<CountedNodePtr> tail;

    // 增加外部引用计数
    static void increase_external_count(std::atomic<CountedNodePtr>& counter,
                                        CountedNodePtr& old_counter) {
        CountedNodePtr new_counter;
        do {
            new_counter = old_counter;
            ++new_counter.external_count;
        } while (!counter.compare_exchange_strong(
                old_counter, new_counter,
                std::memory_order_acquire,
                std::memory_order_relaxed));

        old_counter.external_count = new_counter.external_count;
    }

    // 释放外部计数器
    static void free_external_counter(CountedNodePtr& old_node_ptr) {
        Node* const ptr = old_node_ptr.ptr;
        int const count_increase = old_node_ptr.external_count - 2;

        // 增加内部计数（可能为负）
        int old_count = ptr->internal_count.fetch_add(count_increase,
                                                      std::memory_order_acq_rel);
        if (old_count == -count_increase) {
            delete ptr;
        }
    }

public:
    MPMCLockFreeQueue() {
        CountedNodePtr dummy;
        dummy.ptr = new Node();
        dummy.external_count = 1;
        head.store(dummy, std::memory_order_relaxed);
        tail.store(dummy, std::memory_order_relaxed);
    }

    ~MPMCLockFreeQueue() {
        // 析构时清空队列
        while (pop() != nullptr);
        delete head.load(std::memory_order_relaxed).ptr;
    }

    void push(T new_value) {
        // 创建新节点
        std::unique_ptr<T> new_data(new T(std::move(new_value)));
        CountedNodePtr new_next;
        new_next.ptr = new Node();
        new_next.external_count = 1;

        CountedNodePtr old_tail = tail.load(std::memory_order_relaxed);

        for (;;) {
            increase_external_count(tail, old_tail);
            Node* const old_tail_ptr = old_tail.ptr;

            // 尝试将数据存入当前尾节点
            T* old_data = nullptr;
            if (old_tail_ptr->data.compare_exchange_strong(
                    old_data, new_data.get(),
                    std::memory_order_release,
                    std::memory_order_relaxed)) {

                // 设置新节点的next指针
                CountedNodePtr old_next = {0, nullptr};
                if (!old_tail_ptr->next.compare_exchange_strong(
                        old_next, new_next,
                        std::memory_order_release,
                        std::memory_order_relaxed)) {
                    // 其他线程已经设置了next指针
                    delete new_next.ptr;
                }

                // 尝试移动尾指针
                old_tail = new_next;
                if (tail.compare_exchange_strong(
                        old_tail, new_next,
                        std::memory_order_release,
                        std::memory_order_relaxed)) {
                    free_external_counter(old_tail);
                }

                new_data.release(); // 成功转移所有权
                break;
            }

            // CAS失败，帮助其他线程完成操作
            CountedNodePtr old_next = {0, nullptr};
            if (old_tail_ptr->next.compare_exchange_strong(
                    old_next, new_next,
                    std::memory_order_release,
                    std::memory_order_relaxed)) {

                old_tail = new_next;
                if (tail.compare_exchange_strong(
                        old_tail, new_next,
                        std::memory_order_release,
                        std::memory_order_relaxed)) {
                    free_external_counter(old_tail);
                }
                break;
            }

            old_tail_ptr->release_ref();
            old_tail = tail.load(std::memory_order_relaxed);
        }
    }

    std::shared_ptr<T> pop() {
        CountedNodePtr old_head = head.load(std::memory_order_relaxed);
        for (;;) {
            increase_external_count(head, old_head);
            Node* const ptr = old_head.ptr;

            // 检查队列是否为空
            if (ptr == tail.load(std::memory_order_acquire).ptr) {
                ptr->release_ref();
                return nullptr;
            }

            // 读取下一个节点
            CountedNodePtr next = ptr->next.load(std::memory_order_acquire);
            if (head.compare_exchange_strong(
                    old_head, next,
                    std::memory_order_release,
                    std::memory_order_relaxed)) {

                // 获取数据所有权
                T* res = ptr->data.exchange(nullptr, std::memory_order_acquire);
                if (!res) {
                    // 数据已被其他消费者取出
                    continue;
                }

                // 释放外部计数
                free_external_counter(old_head);
                return std::shared_ptr<T>(res);
            }

            ptr->release_ref();
        }
    }
};


void test_SPSC_basic_operations() {
    std::cout << "Testing SPSC basic operations...\n";
    LockFreeQueue<int> queue;

    // 测试空队列
    assert(queue.empty());
    assert(queue.pop() == nullptr);

    // 测试单元素入队出队
    queue.push(42);
    assert(!queue.empty());
    auto val = queue.pop();
    assert(val != nullptr && *val == 42);
    assert(queue.empty());

    // 测试多个元素
    for (int i = 0; i < 100; ++i) {
        queue.push(i);
    }
    for (int i = 0; i < 100; ++i) {
        auto v = queue.pop();
        assert(v != nullptr && *v == i);
    }
    assert(queue.empty());

    std::cout << "SPSC basic operations test passed!\n";
}

void test_SPSC_single_producer_single_consumer() {
    std::cout << "Testing SPSC concurrent scenario...\n";
    LockFreeQueue<int> queue;
    const int NUM_ITEMS = 100;

    // 生产者线程
    auto producer = [&queue]() {
        for (int i = 0; i < NUM_ITEMS; ++i) {
            queue.push(i);
            cout << " producer  " << i << endl;
        }
    };

    // 消费者线程
    std::vector<int> consumed;
    consumed.reserve(NUM_ITEMS);
    auto consumer = [&queue, &consumed]() {
        for (int i = 0; i < NUM_ITEMS; ++i) {
            while (true) {
                auto item = queue.pop();
                if (item) {
                    consumed.push_back(*item);
                    cout << " consumed  " << *item << endl;
                    break;
                }
                std::this_thread::yield();
            }
        }
    };

    std::thread prod_thread(producer);
    std::thread cons_thread(consumer);

    prod_thread.join();
    cons_thread.join();

    // 验证所有元素都被正确处理且顺序正确
    assert(consumed.size() == NUM_ITEMS);
    for (int i = 0; i < NUM_ITEMS; ++i) {
        assert(consumed[i] == i);
    }

    std::cout << "SPSC concurrent test passed!\n";
}

void run_SPSC_tests() {
    test_SPSC_basic_operations();
    test_SPSC_single_producer_single_consumer();
}



void test_concurrent_ops() {
    MPMCLockFreeQueue<int> queue;
    const int THREADS = 4;
    const int OPS = 100;
    std::atomic<int> counter{0};

    auto producer = [&](int id) {
        for (int i = 0; i < OPS; ++i) {
            queue.push(id * OPS + i);
            cout << "----producer---" << endl;
        }
    };

    auto consumer = [&] {
        while (counter.fetch_add(1) < THREADS * OPS) {
            auto val = queue.pop();
            if (val) {
                cout << "---consumer----" << endl;
            } else {
//                std::this_thread::yield();
            }
        }
    };

    std::vector<std::thread> producers;
    std::vector<std::thread> consumers;

    for (int i = 0; i < THREADS; ++i) {
        producers.emplace_back(producer, i);
    }
    for (int i = 0; i < 2; ++i) {
        consumers.emplace_back(consumer);
    }

    for (auto& t : producers) t.join();
    for (auto& t : consumers) t.join();

    std::cout << "Test completed. Queue empty: "
              << (queue.pop() == nullptr) << std::endl;
}


int main(int argc, char * argv[]){
//    run_SPSC_tests();
    test_concurrent_ops();
    return 0;
}